package com.test.netty;

import org.springframework.boot.web.embedded.netty.NettyReactiveWebServerFactory;
import org.springframework.boot.web.server.WebServerFactoryCustomizer;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import reactor.netty.http.server.HttpServer;
import java.util.concurrent.atomic.AtomicInteger;

@Configuration
public class CustomConnectionLimiter {

    private static final int MAX_CONNECTIONS = 120000;
    private final AtomicInteger currentConnections = new AtomicInteger(0);

    @Bean
    public WebServerFactoryCustomizer<NettyReactiveWebServerFactory> connectionLimiterCustomizer() {
        return factory -> factory.addServerCustomizers(httpServer ->
            httpServer
                .doOnConnection(conn -> {
                    if (currentConnections.incrementAndGet() > MAX_CONNECTIONS) {
                        conn.dispose();
                    }
                    conn.onDispose(() -> currentConnections.decrementAndGet());
                })
        );
    }
}
